import sqlite3
import csv

# 数据库连接
conn = sqlite3.connect("E:\\immunedb\\immuno.db")
cursor = conn.cursor()

try:
    specified_hlas = [
        "HLA-A*02:02", "HLA-A*02:03", "HLA-A*02:04", "HLA-A*02:05", "HLA-A*02:06", 
        "HLA-A*02:07", "HLA-A*02:11", "HLA-A*02:12", "HLA-A*02:16", "HLA-A*02:17", 
        "HLA-A*02:19", "HLA-A*02:20", "HLA-A*02:50", "HLA-A*03:01", "HLA-A*11:01", 
        "HLA-A*23:01", "HLA-A*24:03", "HLA-A*24:06", "HLA-A*24:13", "HLA-A*25:01", 
        "HLA-A*26:01", "HLA-A*26:03", "HLA-A*80:01"
    ]
    print(f"共处理 {len(specified_hlas)} 种HLA类型")

    # 步骤1：创建临时表（
    print("\n=== 步骤1：创建临时表 ===")
    cursor.execute("""
        CREATE TEMP TABLE IF NOT EXISTS temp_joined_data (
            hla TEXT,
            tissue TEXT,
            probability REAL,
            avg_exp REAL  -- 基因平均表达量
        );
    """)
    print("临时表创建成功")

    # 步骤2：分批插入数据
    print("\n=== 步骤2：分批插入数据 ===")
    for i, hla in enumerate(specified_hlas, 1):
        # 查询当前HLA的数据量
        cursor.execute(f"""
            SELECT COUNT(*) 
            FROM predictions p
            JOIN peptides pe ON p.peptide_id = pe.peptide_id
            JOIN protein_gene pg ON pe.protein_id = pg.protein_id
            WHERE p.hla = '{hla}';
        """)
        count = cursor.fetchone()[0]
        print(f"- 第{i}/{len(specified_hlas)}种：{hla}，共 {count} 条数据")

        if count > 0:
            # 插入所有数据
            cursor.execute(f"""
                INSERT INTO temp_joined_data (hla, tissue, probability, avg_exp)
                SELECT 
                    p.hla,
                    pg.tissue,  -- 从protein_gene表获取组织
                    p.probability,
                    pg.avg_exp  -- 从protein_gene表获取平均表达量
                FROM predictions p
                JOIN peptides pe ON p.peptide_id = pe.peptide_id
                JOIN protein_gene pg ON pe.protein_id = pg.protein_id
                WHERE p.hla = '{hla}';
            """)
            conn.commit()
            print(f"  已插入临时表\n")
        else:
            print(f"  无数据，跳过\n")

    # 步骤3：计算全局平均值阈值
    print("\n=== 步骤3：计算avg_exp平均值（仅prob>0.5的肽段） ===")
    cursor.execute("""
        SELECT 
            ROUND(
                CAST(
                    SUM(CASE WHEN probability > 0.5 THEN avg_exp ELSE 0 END)
                AS REAL)
                / MAX(SUM(CASE WHEN probability > 0.5 THEN 1 ELSE 0 END), 1),
                4
            ) AS avg_exp_mean
        FROM temp_joined_data;
    """)
    avg_exp_mean = cursor.fetchone()[0]
    if avg_exp_mean == 0 and not any(cursor.fetchall()):
        raise ValueError("临时表中没有probability>0.5的数据，无法计算avg_exp平均值")
    print(f"所有记录的avg_exp平均值（仅prob>0.5）为：{avg_exp_mean:.4f}")

    # 步骤4：按HLA+组织统计并分类
    print("\n=== 步骤4：统计并分类结果 ===")
    output_path = "E:/immunedb/RNA_heatmap.csv"
    with open(output_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        # 
        writer.writerow(["HLA", "组织类型", "概率>0.5的数量", "概率<0.5的数量", 
                        "avg_exp平均值（仅prob>0.5）", "表达量分类(high/low/medium)"])

        for hla in specified_hlas:
            # 获取当前HLA对应的所有组织
            cursor.execute(f"""
                SELECT DISTINCT tissue 
                FROM temp_joined_data 
                WHERE hla = '{hla}';
            """)
            tissues = [t[0] for t in cursor.fetchall()]

            for tissue in tissues:
                # SUM(CASE)求和 + MAX避免除0 + ROUND保留精度
                cursor.execute(f"""
                    SELECT 
                        '{hla}' AS hla,
                        '{tissue}' AS tissue,
                        SUM(CASE WHEN probability > 0.5 THEN 1 ELSE 0 END) AS count_gt_0_5,
                        SUM(CASE WHEN probability < 0.5 THEN 1 ELSE 0 END) AS count_lt_0_5,
                        ROUND(
                            CAST(
                                SUM(CASE WHEN probability > 0.5 THEN avg_exp ELSE 0 END)
                            AS REAL)
                            / MAX(SUM(CASE WHEN probability > 0.5 THEN 1 ELSE 0 END), 1),
                            4
                        ) AS tissue_avg_exp
                    FROM temp_joined_data
                    WHERE hla = '{hla}' AND tissue = '{tissue}';
                """)
                result = cursor.fetchone()
                if not result:
                    continue  # 跳过无数据的组合

                # 解析结果（result已包含hla、tissue、count_gt_0_5、count_lt_0_5、tissue_avg_exp）
                hla_val, tissue_val, count_gt_0_5, count_lt_0_5, tissue_avg = result
                # 处理None值
                count_gt_0_5 = count_gt_0_5 or 0
                count_lt_0_5 = count_lt_0_5 or 0
                tissue_avg = tissue_avg or 0

                # 按全局平均值判断分类
                if tissue_avg > avg_exp_mean:
                    exp_class = "high"
                elif tissue_avg < avg_exp_mean:
                    exp_class = "low"
                else:
                    exp_class = "medium"

                # 写入CSV
                writer.writerow([
                    hla_val,
                    tissue_val,
                    count_gt_0_5,
                    count_lt_0_5,
                    tissue_avg,
                    exp_class
                ])
                print(f"已处理：{hla_val} + {tissue_val} → 平均表达量：{tissue_avg:.4f}，分类：{exp_class}")

    print(f"\n结果已保存到：{output_path}")

except sqlite3.Error as e:
    print(f"\n数据库错误: {e}")
except Exception as e:
    print(f"\n其他错误: {e}")
finally:
    cursor.close()
    conn.close()
    print("\n程序结束，数据库连接已关闭")

